# 画框和可视化模块
import cv2
import numpy as np
from typing import List, Dict, Any, Tuple


class Visualizer:
    """可视化工具类"""

    def __init__(self, vis_config: Dict[str, Any]):
        """
        初始化可视化器
        Args:
            vis_config: 可视化配置
        """
        self.person_color = tuple(vis_config['person_color'])
        self.other_color = tuple(vis_config['other_color'])
        self.line_thickness = vis_config['line_thickness']
        self.font_scale = vis_config['font_scale']
        self.font = cv2.FONT_HERSHEY_SIMPLEX

    def draw_detections(self, frame: np.ndarray,
                        detections: List[Dict[str, Any]]) -> np.ndarray:
        """
        在帧上绘制检测结果
        Args:
            frame: 输入图像帧
            detections: 检测结果列表
        Returns:
            绘制后的图像帧
        """
        frame_copy = frame.copy()

        for detection in detections:
            bbox = detection['bbox']
            class_name = detection['class_name']
            confidence = detection['confidence']
            class_id = detection['class_id']

            x1, y1, x2, y2 = bbox

            # 根据类别选择颜色
            if class_id == 0:  # person
                color = self.person_color
            else:
                color = self.other_color

            # 绘制边界框
            cv2.rectangle(frame_copy, (x1, y1), (x2, y2), color, self.line_thickness)

            # 准备标签文本
            label = f"{class_name}: {confidence:.2f}"

            # 获取文本大小
            (text_width, text_height), baseline = cv2.getTextSize(
                label, self.font, self.font_scale, 1
            )

            # 绘制标签背景
            cv2.rectangle(
                frame_copy,
                (x1, y1 - text_height - baseline - 10),
                (x1 + text_width, y1),
                color,
                -1
            )

            # 绘制标签文本
            cv2.putText(
                frame_copy,
                label,
                (x1, y1 - baseline - 5),
                self.font,
                self.font_scale,
                (255, 255, 255),  # 白色文字
                1,
                cv2.LINE_AA
            )

        return frame_copy

    def add_info_overlay(self, frame: np.ndarray, info: Dict[str, Any]) -> np.ndarray:
        """
        添加信息覆盖层
        Args:
            frame: 输入图像帧
            info: 信息字典(如FPS等)
        Returns:
            添加信息后的图像帧
        """
        frame_copy = frame.copy()

        # 添加FPS信息
        if 'fps' in info:
            fps_text = f"FPS: {info['fps']:.1f}"
            cv2.putText(
                frame_copy,
                fps_text,
                (10, 30),
                self.font,
                0.7,
                (0, 255, 0),
                2,
                cv2.LINE_AA
            )

        # 添加检测数量信息
        if 'detection_count' in info:
            count_text = f"Detections: {info['detection_count']}"
            cv2.putText(
                frame_copy,
                count_text,
                (10, 60),
                self.font,
                0.7,
                (0, 255, 0),
                2,
                cv2.LINE_AA
            )

        return frame_copy
